﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;

namespace ImageProcessingLibrary
{
    internal unsafe class ConnectedComponent
    {
        #region Fields
        private int[,] connectedComponentsLabels;
        private readonly int height;
        private readonly int width;
        private int[] regionSizeByLabel; //Background is label 0
        #endregion

        #region Constructor
        internal ConnectedComponent(ImageProcessor binarizedImage)
        {
            lock (this)
            {
                this.height = binarizedImage.processedData.Height;
                this.width = binarizedImage.processedData.Width;
                FindConnectedComponents(binarizedImage);
            }
        }

        // Used only by constructor
        private void FindConnectedComponents(ImageProcessor binarizedImage)
        {
            // Three-pass algorithm. Variant of: http://en.wikipedia.org/wiki/Connected_Component_Labeling#Two-pass
            // But one pass is introduced so that we may have a fixed-sized array to represent DisjointSets.

            int[,] labels = new int[binarizedImage.ProcessedBitmap.Height, binarizedImage.ProcessedBitmap.Width];
            labels.Initialize(); // Init to 0 (background)
            int nrOfLabels = 1; //Background is label 0

            // First pass
            // Label all pixels. Count number of labels.
            #region First pass
            // First element
            byte* row0 = (byte*)binarizedImage.processedData.Scan0;
            if (row0[0] != 0)
            {
                labels[0, 0] = nrOfLabels;
                nrOfLabels++;
            }
            // First Row:
            for (int j = 1; j < width; j++)
            {
                byte pixelColor = (byte)(row0[j * ImageProcessor.PIXEL_SIZE]);
                if (pixelColor != 0 && labels[0, (j - 1)] != 0)
                {
                    labels[0, j] = labels[0, (j - 1)];
                }
                else if (pixelColor != 0)
                {
                    labels[0, j] = nrOfLabels;
                    nrOfLabels++;
                }
            }

            //First Column:
            for (int i = 1; i < height; i++)
            {
                byte* pRow = (byte*)binarizedImage.processedData.Scan0 + (i * binarizedImage.processedData.Stride);
                byte pixelColor = (byte)(pRow[0]);
                if (pixelColor != 0 && labels[(i - 1), 0] != 0)
                {
                    labels[i, 0] = labels[(i - 1), 0];
                }
                else if (pixelColor != 0)
                {
                    labels[i, 0] = nrOfLabels;
                    nrOfLabels++;
                }
            }

            //The rest:
            for (int i = 1; i < height; i++)
            {
                byte* pRow = (byte*)binarizedImage.processedData.Scan0 + (i * binarizedImage.processedData.Stride);
                for (int j = 1; j < width; j++)
                {
                    byte pixelColor = (byte)(pRow[j * ImageProcessor.PIXEL_SIZE]);
                    int westLabel = labels[i, j - 1];
                    int northLabel = labels[i - 1, j];

                    if (pixelColor != 0 && westLabel != 0)
                    {
                        labels[i, j] = westLabel;
                    }
                    else if (pixelColor != 0 && northLabel != 0)
                    {
                        labels[i, j] = northLabel;
                    }
                    else if (pixelColor != 0)
                    {
                        labels[i, j] = nrOfLabels;
                        nrOfLabels++;
                    }
                }
            }
            #endregion

            // Second pass
            // Create DisjointSets with same size as number of lables. Union all connected sets.
            #region Second pass
            DisjointSets disjointSets = new DisjointSets(nrOfLabels); //Label 0 is backgound

            for (int i = 1; i < height; i++)
            {
                for (int j = 1; j < width; j++)
                {
                    int currentLabel = labels[i, j];
                    int westLabel = labels[i, j - 1];
                    int northLabel = labels[i - 1, j];

                    if (currentLabel != 0 && westLabel != 0)
                    {
                        disjointSets.Union(currentLabel, westLabel);
                    }
                    if (currentLabel != 0 && northLabel != 0)
                    {
                        disjointSets.Union(currentLabel, northLabel);
                    }
                }
            }
            #endregion

            // Third pass
            // Relabel all. Count Region size.
            #region Third pass
            int[] regionSizeByLabel = new int[nrOfLabels];
            regionSizeByLabel.Initialize();

            for (int i = 0; i < height; i++)
            {
                for (int j = 0; j < width; j++)
                {
                    int currentLabel = labels[i, j];
                    if (currentLabel != 0) // Don't count background
                    {
                        int jointSetLabel = disjointSets.Find(currentLabel);
                        regionSizeByLabel[jointSetLabel]++;
                        labels[i, j] = jointSetLabel;
                    }
                }
            }
            #endregion

            this.regionSizeByLabel = regionSizeByLabel;
            this.connectedComponentsLabels = labels;
        }
        #endregion

        #region Internal methods
        internal void RemoveSmallComponents(int threshold)
        {
            lock (this)
            {
                List<int> labelsToRemove = new List<int>();
                for(int i = 1; i < this.regionSizeByLabel.Length; i++)
                {
                    if (this.regionSizeByLabel[i] < threshold)
                    {
                        labelsToRemove.Add(i);
                    }
                }

                this.RemoveComponents(labelsToRemove);
            }
        }

        internal void RemoveNonFaces(double widthVsHeightLow, double widthVsHeightHigh)
        {
            lock (this)
            {
                double[] shapeByLabel = this.ShapeOfRegions();
                List<int> lablesToRemove = new List<int>();

                for (int i = 1; i < shapeByLabel.Length; i++)
                {
                    if (shapeByLabel[i] < widthVsHeightLow || shapeByLabel[i] > widthVsHeightHigh)
                    {
                        lablesToRemove.Add(i);
                    }
                }

                this.RemoveComponents(lablesToRemove);
            }
        }

        internal void KeepOnlyConnectedPixels(ImageProcessor image)
        {
            lock (this)
            {
                for (int i = 0; i < height; i++)
                {
                    byte* pRow = (byte*)image.processedData.Scan0 + (i * image.processedData.Stride);
                    for (int j = 0; j < width; j++)
                    {
                        if (this.connectedComponentsLabels[i, j] == 0) // Only background
                        {
                            pRow[j * ImageProcessor.PIXEL_SIZE] = 0; // B
                            pRow[j * ImageProcessor.PIXEL_SIZE + 1] = 0; //G
                            pRow[j * ImageProcessor.PIXEL_SIZE + 2] = 0; //R
                        }
                    }
                }
            }
        }
        #endregion

        #region Private methods
        private double[] ShapeOfRegions()
        {
            int l = this.regionSizeByLabel.Length;

            // Find corners
            int[] lefts = new int[l];
            int[] rights = new int[l];
            int[] tops = new int[l];
            int[] bottoms = new int[l];

            for (int i = 1; i < l; i++)
            {
                lefts[i] = width;
                rights[i] = 0;
                tops[i] = height;
                bottoms[i] = 0;
            }

            for (int i = 0; i < height; i++)
            {
                for (int j = 0; j < width; j++)
                {
                    int label = this.connectedComponentsLabels[i, j];
                    if (label != 0)
                    {
                        if (lefts[label] > j) // Find min
                        {
                            lefts[label] = j;
                        }
                        if (rights[label] < j) // Find max
                        {
                            rights[label] = j;
                        }
                        if (tops[label] > i) // Find min
                        {
                            tops[label] = i;
                        }
                        if (bottoms[label] < i) // Find max
                        {
                            bottoms[label] = i;
                        }
                    }
                }
            }

            // Calculate proportions
            double[] shapeByLabel = new double[l];
            shapeByLabel.Initialize();

            for (int i = 1; i < l; i++)
            {
                double w = (double)(rights[i] - lefts[i]);
                double h = (double)(bottoms[i] - tops[i]);

                if (h == 0.0)
                {
                    shapeByLabel[i] = 1000; // A large number
                }
                else
                {
                    shapeByLabel[i] = w/h;
                }
            }

            return shapeByLabel;
        }

        private void RemoveComponents(List<int> labelsToRemove)
        {
            int[] newRegionSizeByLabel = new int[this.regionSizeByLabel.Length - labelsToRemove.Count];
            newRegionSizeByLabel.Initialize();

            // Create a lookuptable for the new labels
            int[] reLabelTable = new int[this.regionSizeByLabel.Length];
            reLabelTable.Initialize();
            for (int iOldID = 1; iOldID < reLabelTable.Length; iOldID++)
            {
                if (labelsToRemove.Contains(iOldID))
                {
                    reLabelTable[iOldID] = 0; //Background
                }
                else
                {
                    List<int> smallerRemovedLabels = labelsToRemove.FindAll(Label => Label < iOldID);
                    Debug.Assert(smallerRemovedLabels.Count >= 0);
                    int newID = iOldID - smallerRemovedLabels.Count;
                    reLabelTable[iOldID] = newID;

                    //Transfer region size to new label
                    newRegionSizeByLabel[newID] = this.regionSizeByLabel[iOldID];
                }
            }

            //Relabel all
            for (int i = 0; i < height; i++)
            {
                for (int j = 0; j < width; j++)
                {
                    int oldLabel = this.connectedComponentsLabels[i, j];
                    int newLabel = reLabelTable[oldLabel];
                    this.connectedComponentsLabels[i, j] = newLabel;
                }
            }

            //Update Sizes
            this.regionSizeByLabel = newRegionSizeByLabel;
        }
        #endregion

        #region Properties
        internal int NumberOfComponents
        {
            get
            {
                if (this.regionSizeByLabel.Length < 1)
                    return 0;
                else
                {
                    return this.regionSizeByLabel.Length - 1; //Don't count background
                }
            }
        }
        #endregion
    }
}
